import json
import os
import pandas as pd
from selenium import webdriver
import requests
import random
import time
from typing import Union
from selenium.common.exceptions import WebDriverException


# 封装获取代理IP的函数
def get_proxies(proxies: Union[str, list]):
    """
    :param proxies: 字符串格式代理：'http://127.0.0.1:8080'
                    或列表形式多个代理：[
                    'http://127.0.0.1:8080',
                    'https://127.0.0.1:8080',
                    'http://127.0.0.1:8888',
                    'https://127.0.0.1:8888']
    :return: 构建的代理字典
    """
    if isinstance(proxies, str):
        proxies = {'http': proxies, 'https': proxies}
    elif isinstance(proxies, list):
        random_proxy = random.choice(proxies)
        proxies = {'http': random_proxy, 'https': random_proxy}
    else:
        print("Invalid input format for proxies. It should be either a string or a list of proxies.")
        proxies = None
    return proxies


# 封装selenium创建浏览器对象的函数
def create_chrome_driver(*, headless=False, download_path=None, proxies=None):
    """
    :param headless: 是否启用无头浏览器，默认为False不启用
    :param download_path: 设置默认下载路径，默认为None不设置
    :param proxies: 设置代理，默认为None不设置
    :return: 浏览器对象
    """
    try:
        options = webdriver.ChromeOptions()

        # 处理SSL证书错误问题
        options.add_argument('--ignore-certificate-errors')
        options.add_argument('--ignore-ssl-errors')
        # 禁用信息栏
        # options.add_argument("--disable-infobars")

        if headless:
            # 无头浏览器
            options.add_argument('--headless')
        if download_path:
            # 设置默认下载路径
            prefs = {"download.default_directory": download_path}
            options.add_experimental_option("prefs", prefs)
        if proxies:
            # 设置代理
            proxy_str = f"{proxies['http']}"
            options.add_argument(f'--proxy-server={proxy_str}')
        options.add_experimental_option('excludeSwitches', ['enable-automation', 'enable-logging'])
        options.add_experimental_option('useAutomationExtension', False)
        browser = webdriver.Chrome(options=options)
        browser.execute_cdp_cmd(
            'Page.addScriptToEvaluateOnNewDocument',
            {'source': 'Object.defineProperty(navigator, "webdriver", {get: () => undefined})'}
        )
        return browser
    except WebDriverException as e:
        print(f"An error occurred：{e}")
        return None


# 封装浏览器对象获取并保存cookies的函数
def save_cookies(url, sleep_time, cookie_file=None):
    """
    :param url: 请求地址
    :param cookie_file: cookie存储文件路径
    :param sleep_time: 等待时间
    :return: cookie信息
    """
    browser = webdriver.Chrome()
    browser.get(url)
    print(f"The page is loaded, please log in manually and wait {sleep_time} seconds")
    time.sleep(sleep_time)
    # 获取cookie信息并保存到本地
    cookie_json = browser.get_cookies()
    if cookie_file:
        with open(cookie_file, 'w', encoding='utf-8') as w_file:
            json.dump(cookie_json, w_file)
    browser.close()
    return cookie_json


# 封装浏览器对象添加cookies的函数
def add_cookies(browser, cookie_data: Union[str, list]):
    """
    :param browser: 浏览器对象
    :param cookie_data: cookie文件路径或者直接的cookie数据
    :return:
    """
    if isinstance(cookie_data, str):  # 如果输入是文件路径
        with open(cookie_data, 'r') as f:
            cookies_list = json.load(f)
    elif isinstance(cookie_data, list):  # 如果输入是直接的cookie数据
        cookies_list = cookie_data
    else:
        raise ValueError("Invalid input format for cookie_data. It should be either a file path or a list of cookies.")

    for cookie_dict in cookies_list:
        if cookie_dict.get('secure'):
            browser.add_cookie(cookie_dict)


# 封装发送请求的函数，按照返回类型返回响应内容
def send_request(url, method='GET', return_type='text', headers=None, params=None,
                 data=None, json_data=None, verify=True, proxy=None):
    """
    :param url: 请求地址
    :param method: 请求方式，默认为GET
    :param return_type: 响应内容类型，默认为text，可选json、content
    :param headers: 请求头
    :param params: 请求参数
    :param data: 请求体
    :param json_data: 请求体json格式
    :param verify: 请求是否验证SSL证书，默认为True
    :param proxy: 请求代理，默认为None
    :return: 响应内容
    """
    if headers is None:
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
                          "Chrome/123.0.0.0 Safari/537.36"
        }
    try:
        if method.upper() == 'GET':
            response = requests.get(url, headers=headers, params=params, verify=verify, proxies=proxy)
        elif method.upper() == 'POST':
            if json_data:
                response = requests.post(url, headers=headers, json=json_data, verify=verify, proxies=proxy)
            else:
                response = requests.post(url, headers=headers, data=data, verify=verify, proxies=proxy)
        else:
            print(f"Unsupported request method: {method}")
            return None
        if response.status_code == 200:
            if return_type == 'text':
                return response.text
            elif return_type == 'json':
                return response.json()
            elif return_type == 'content':
                return response.content
            else:
                return None
        else:
            print(f"Request failed, status code: {response.status_code}")
            return None
    except requests.exceptions.RequestException as e:
        print(f"Request exception: {e}")
        return None


# 封装逐条数据写入CSV文件的函数
def append_to_csv(data_list, path, columns):
    """
    :param data_list: 要存储的单条数据列表
    :param path: 数据存储CSV文件路径
    :param columns: CSV文件列名列表
    :return:
    """
    # 将数据转换为DataFrame
    data_df = pd.DataFrame([data_list], columns=columns)
    # 检查文件是否存在来决定是否写入表头
    header = not os.path.isfile(path)
    # 追加数据到CSV文件，不写入索引
    data_df.to_csv(path, mode='a', index=False, header=header)
